import time

import mujoco.viewer
import mujoco
import numpy as np
import torch
import yaml

from my_lcm import Request, Response
import lcm
lcm = lcm.LCM()


#接收action
def my_handler(channel, data):
    global action_rl
    msg = Response.decode(data)
    #print("q_exp_rl:",msg.q_exp)
    action_rl=msg.q_exp


def get_gravity_orientation(quaternion):
    qw = quaternion[0]
    qx = quaternion[1]
    qy = quaternion[2]
    qz = quaternion[3]

    gravity_orientation = np.zeros(3)

    gravity_orientation[0] = 2 * (-qz * qx + qw * qy)
    gravity_orientation[1] = -2 * (qz * qy + qw * qx)
    gravity_orientation[2] = 1 - 2 * (qw * qw + qz * qz)

    return gravity_orientation


def pd_control(target_q, q, kp, target_dq, dq, kd):
    """Calculates torques from position commands"""
    return (target_q - q) * kp + (target_dq - dq) * kd


if __name__ == "__main__":

    policy_path = "/home/saw/C++/sim2sim/model/policy_5.pt" 
    xml_path = "/home/saw/C++/sim2sim/xdog/mjcf1/world.xml"

    #训练相关参数
    num_actions = 12
    num_obs = 45
    default_angles = [0.0,-1.0,2.0,0.0,-1.0,2.0,0.0,-1.0,2.0,0.0,-1.0,2.0]
    action_scale=0.25
    default_dof_pos = [0.0,-1.0,2.0,0.0,-1.0,2.0,0.0,-1.0,2.0,0.0,-1.0,2.0]
    
    action_rl=default_dof_pos #初始化action_rl

    #仿真相关变量
    simulation_duration = 60.0
    simulation_dt = 0.002
    control_decimation = 10

    #kp kd
    kp = [20.0, 20.0, 20.0, 20.0, 20.0,  20.0, 20.0, 20.0, 20.0, 20.0, 20.0,20.0]
    kd = [0.5, 0.5, 0.5, 0.5, 0.5,  0.5, 0.5, 0.5, 0.5, 0.5,0.5,0.5]
    kps = np.array(kp, dtype=np.float32)
    kds = np.array(kd, dtype=np.float32)

    # define context variables
    action = np.zeros(num_actions, dtype=np.float32)
    target_dof_pos = default_angles.copy()
    obs = np.zeros(num_obs, dtype=np.float32)

    #计数变量
    counter = 0

    # Load robot model
    m = mujoco.MjModel.from_xml_path(xml_path)
    d = mujoco.MjData(m)
    m.opt.timestep = simulation_dt

    # 加载policy
    policy = torch.jit.load(policy_path)

    with mujoco.viewer.launch_passive(m, d) as viewer:
        # Close the viewer automatically after simulation_duration wall-seconds.

        start = time.time()
        while viewer.is_running() and time.time() - start < simulation_duration:
            step_start = time.time()
            tau = pd_control(target_dof_pos, d.qpos[7:], kps, np.zeros_like(kds), d.qvel[6:], kds)
            d.ctrl[:] = tau
            # mj_step can be replaced with code that also evaluates
            # a policy and applies a control signal before stepping the physics.
            mujoco.mj_step(m, d)

            #订阅action
            subscription = lcm.subscribe("LCM_ACTION", my_handler)

            counter += 1
            if counter % control_decimation == 0:

                #创建observation
                qj = d.qpos[7:] #关节角度（12维）
                dqj = d.qvel[6:] #关节角速度（12维）
                quat = d.qpos[3:7] #角度四元数
                omega = d.qvel[3:6] #角速度（3 维）
                base_vel = d.qvel[:3] #线速度（3 维）
                gravity_orientation = get_gravity_orientation(quat) #重力投影

                
                msg = Request()
                msg.omega = omega
                msg.gravity_orientation = gravity_orientation
                #msg.command = cmd
                msg.q = qj
                msg.dq = dqj
                
                lcm.publish("LCM_OBS", msg.encode())
                lcm.handle()

                # 确保转为numpy数组
                action = np.array(action_rl)  

                target_dof_pos = action * action_scale + default_angles

            # Pick up changes to the physics state, apply perturbations, update options from GUI.
            viewer.sync()

            # Rudimentary time keeping, will drift relative to wall clock.
            time_until_next_step = m.opt.timestep - (time.time() - step_start)
            if time_until_next_step > 0:
                time.sleep(time_until_next_step)
